from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch


class Translator:
    def __init__(self, model_name="facebook/nllb-200-distilled-600M", target_lang="zho_Hans"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(
            "cuda" if torch.cuda.is_available() else "cpu")
        self.target_lang_id = self.tokenizer.convert_tokens_to_ids(target_lang)

    def translate(self, text: str) -> str:
        if not text.strip(): return ""
        inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)
        output = self.model.generate(**inputs, max_length=128, num_beams=2, early_stopping=True,
                                     forced_bos_token_id=self.target_lang_id)
        return self.tokenizer.decode(output[0], skip_special_tokens=True)
